import torch
from torchvision import transforms
from PIL import Image
import torch.nn as nn
import numpy as np
import albumentations as aug
import albumentations.pytorch as aug_torch
from train_eval_fund_4 import DeepSurModel

# 配置参数
MODEL_PATH = None  # 替换为你的模型路径
CLASS_NAMES = ['正常_N', '糖尿病_D', '青光眼_G', '白内障_C', 'AMD_A', '高血压_H', '近视_M', '其他_O']
IMG_SIZE = 512  # 根据模型输入尺寸修改
THRESHOLD = 0.5  # 分类阈值

# 加载模型
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = DeepSurModel().to(device)
if MODEL_PATH != None:
    model.load_state_dict(torch.load(MODEL_PATH, map_location=device))
model.eval()

# 图像预处理
transform = aug.Compose([
    aug.SmallestMaxSize(
        max_size=IMG_SIZE, always_apply=True),
    aug.CenterCrop(IMG_SIZE, IMG_SIZE,
                    always_apply=True),
    aug.ToFloat(always_apply=True),
    aug_torch.ToTensorV2(),
])

def predict_image(image_path):
    # 加载并预处理图像
    img = Image.open(image_path).convert('RGB')
    img = np.array(img)
    img_tensor = transform(image=img)['image'].unsqueeze(0).to(device)
    
    # 预测
    with torch.no_grad():
        bi_logits, disease_logits = model(img_tensor)
        # print(bi_logits)
        # print(disease_logits)
        logits = torch.cat((bi_logits, disease_logits), dim=1)
        outputs = torch.sigmoid(logits)
        preds = outputs.cpu().numpy()[0]
    
    # 获取结果
    results = {}
    for i, prob in enumerate(preds):
        if prob > THRESHOLD:
            results[CLASS_NAMES[i]] = float(prob)
    
    return results

if __name__ == '__main__':
    # image_path = input()
    predictions = predict_image("/home/zhangyichi/dataset/OIA-ODIR/On-site Test Set/Images/198_left.jpg")
    
    if predictions:
        print("\n分类结果：")
        for cls, prob in predictions.items():
            print(f"{cls}: {prob:.4f}")
    else:
        print("未检测到任何分类结果")